package oktx.feign.loadbalance;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import oktx.participant.JtaParticipant;
import oktx.RemoteType;
import oktx.TransactionContext;

import java.util.Set;
import java.util.stream.Collectors;

/**
 * Created by huangdachao on 2018/5/30 16:06.
 */
public class RestTxLoadBalanceRuleWrapper extends AbstractLoadBalancerRule {
    private IRule rule;

    public RestTxLoadBalanceRuleWrapper(IRule rule) {
        this.rule = rule;
    }

    @Override
    public Server choose(Object key) {
        TransactionContext context = TransactionContext.getContext();
        if (context != null) {
            Set<String> visited = context.getParticipants().stream().filter(p -> p instanceof JtaParticipant && p.getType() == RemoteType.dubbo)
                .map(p -> p.getIp() + p.getPort()).collect(Collectors.toSet());
            Server server = getLoadBalancer().getAllServers().stream().filter(s -> visited.contains(s.getHost() + s.getPort()))
                .findFirst().orElse(null);
            if (server != null) {
                return server;
            }
        }
        return this.rule.choose(key);
    }

    @Override
    public void initWithNiwsConfig(IClientConfig clientConfig) { }
}
